
from model import CNN_BE,UNet
from dataset import Trainset_PDE_generated,Trainset_noise_generated


import time
import math
import os
import torch
import numpy as np
import torch.optim as optim
from torch import linalg
from torch.optim.lr_scheduler import StepLR
from loss import loss_be, operatorx, operatory
from utils import save_checkpoints
from numerical_solution import Numerical_solution

import matplotlib.cm as cm
import matplotlib.pyplot as plt

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
class Tester():
    def __init__(self):

        self.height = 50
        self.width =50
        self.intial_snap = 1
        self.mid_channels = 16
        self.deltaT = 0.001
        self.domain=[-1,1,-1,1]
        self.type = 'BE'
        self.TotalSnaps = 10
        self.TotalTrainingBatch = 1

        self.device = 'cuda'

        self.model_type = 'CNN'
        # self.model_type = 'UNet'

        if self.model_type == 'CNN':
            self.BECNN = CNN_BE(self.mid_channels)
        elif self.model_type == 'UNet':
            self.BECNN = UNet(self.intial_snap, 1)

        self.PDETrainSet = Trainset_noise_generated(self.intial_snap, self.height, self.width, self.domain)
        self.Snap_Training = self.PDETrainSet()

        self.PDETrainSet0 = Trainset_PDE_generated(self.intial_snap, self.height, self.width, self.domain)
        self.Snap_Training0 = self.PDETrainSet0()


        print(self.BECNN)

        if self.device =='cuda':
            self.BECNN = self.BECNN.to(self.device)
            self.Snap_Training_LB = self.Snap_Training0.to(self.device)[:,:,1:,0:-1]
            self.Snap_Training_LB0 = self.Snap_Training0.to(self.device)[:,:,1:,0:-1]
        self.model_name = f'CNNBE_deltaT{self.deltaT}_{self.type}_{self.TotalTrainingBatch}Batch{self.TotalSnaps}Snaps_{self.mid_channels}Channels'
        self.Hx = operatorx(self.Snap_Training_LB)
        self.Hy = operatory(self.Snap_Training_LB)

    def test(self,TimeSteps):
        self.BECNN.eval()
        best_model = torch.load(f'checkpoints/{self.model_name}/checkpoint.pth.tar')
        self.BECNN.load_state_dict(best_model['state_dict'])
        Snap_init =torch.clone(self.Snap_Training_LB0)
        Snap_init1 =torch.clone(self.Snap_Training_LB0)

        TimeSteps = 20
        n=TimeSteps

        for i in range(n):
            Snap_Test_Next = self.BECNN(Snap_init)
            Snap_init = Snap_Test_Next
            snap_n = self.Snap_Training0
            snap_n1 = snap_n[:, :, 1:, 0:-1]
            E = math.exp(-5 * (math.pi) * (math.pi) * i * self.deltaT)
            Exact = (E * snap_n1).to(self.device)
            error = linalg.norm(Snap_init - Exact)
            print(f'{i}, error:{error}')

        Snap_Test_Next_Out = torch.zeros_like(self.Snap_Training0)
        Snap_Test_Next_Out[:, :, 1:, 0:-1] = Snap_Test_Next
        Snap_Test_Next_Out[:, :, 0, 0:-1] = Snap_Test_Next[:, :, self.height - 2, :]
        Snap_Test_Next_Out[:, :, :, self.width - 1] = Snap_Test_Next_Out[:, :, :, 0]

        for j in range(n):
            Snap_Test_New = Numerical_solution(Snap_init1, self.deltaT, self.domain, self.Hx, self.Hy, self.type)
            Snap_init1 = Snap_Test_New
        Snap_Numerical_Out = torch.zeros_like(self.Snap_Training0)
        Snap_Numerical_Out[:, :, 1:, 0:-1] = Snap_Test_New
        Snap_Numerical_Out[:, :, 0, 0:-1] = Snap_Test_New[:, :, self.height - 2, :]
        Snap_Numerical_Out[:, :, :, self.width - 1] = Snap_Numerical_Out[:, :, :, 0]



        E = math.exp(-5 * (math.pi) * (math.pi) * n * self.deltaT)



        snap_initial=self.Snap_Training0[0][0].cpu().detach().numpy()
        snap_next=Snap_Test_Next_Out[0][0].cpu().detach().numpy()
        snap_new_Num = Snap_Numerical_Out[0][0].cpu().detach().numpy()
        a1 = E * snap_initial
        predict_error=np.linalg.norm(snap_next - a1)/np.linalg.norm(a1)
        Numerical_error_Num = np.linalg.norm(snap_new_Num - a1) / np.linalg.norm(a1)


        print(predict_error)
        print(Numerical_error_Num)

        fig, ax =plt.subplots()
        im = ax.imshow(snap_initial, interpolation='bilinear',cmap=cm.RdYlGn)
        plt.title('Initial Snap')
        plt.colorbar(im)
        plt.show()

        fig, ax1 =plt.subplots()
        im1 = ax1.imshow(snap_next, interpolation='bilinear',cmap=cm.RdYlGn)
        plt.title(f'Predicted Snap at T={n*self.deltaT}')
        plt.colorbar(im1)
        plt.show()


        fig, ax3 =plt.subplots()
        im3 = ax3.imshow(a1, interpolation='bilinear',cmap=cm.RdYlGn)
        plt.colorbar(im3)
        plt.title(f'Exact Snap at T={n*self.deltaT}')
        plt.show()

        fig, ax4 =plt.subplots()
        im4 = ax4.imshow(snap_next - a1, interpolation='bilinear',cmap=cm.RdYlGn)
        plt.colorbar(im4)
        plt.title(f'Error at T={n*self.deltaT}, RL2={predict_error}')
        plt.show()

        fig, ax5 = plt.subplots()
        im5 = ax5.imshow(snap_new_Num-a1, interpolation='bilinear', cmap=cm.RdYlGn)
        plt.title(f'Numerical error at T={n * self.deltaT}, RL2={Numerical_error_Num}')
        plt.colorbar(im5)
        plt.show()



if __name__ == '__main__':
    tester = Tester()
    tester.test(50)





